import os, io, base64, math
import json
import cv2
import time
import numpy as np
import pandas as pd
from pathlib import Path
from PIL import Image, ImageDraw
import glob

### json to mask
###################################################################################
# image to array
def img_b64_to_arr(img_b64):
    f = io.BytesIO()
    f.write(base64.b64decode(img_b64))
    img_arr = np.array(Image.open(f))
    return img_arr


# Convert shape to mask
def shape_to_mask(img_shape, points, shape_type=None, line_width=10, point_size=5):
    
    mask = np.zeros(img_shape[:2], dtype=np.uint8) # generate numpy (H * W)

    mask = Image.fromarray(mask) # convert nupmy to PIL

    draw = ImageDraw.Draw(mask)
    
    xy = [tuple(point) for point in points] # generate (x,y)
    
    if shape_type == 'circle':
        assert len(xy) == 2, 'Shape of shape_type=circle must have 2 points'
        (cx, cy), (px, py) = xy
        d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2)
        draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1) 
    elif shape_type == 'rectangle':
        assert len(xy) == 2, 'Shape of shape_type=rectangle must have 2 points'
        draw.rectangle(xy, outline=1, fill=1)
    elif shape_type == 'line':
        assert len(xy) == 2, 'Shape of shape_type=line must have 2 points'
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == 'linestrip':
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == 'point':
        assert len(xy) == 1, 'Shape of shape_type=point must have 1 points'
        cx, cy = xy[0]
        r = point_size
        draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=1, fill=1)
    else:
        mask = np.zeros(img_shape[:2], dtype=np.uint8) # generate numpy (H * W)
        cv2.drawContours(mask, [np.array(points, dtype=np.int32)], 0, 255, -1)
        mask = np.array(mask, dtype=bool)
        return mask
        # if len(xy) <= 2:
        #     # 'Polygon must have points more than 2'
        #     pass
        # else:
        #     draw.polygon(xy=xy, outline=1, fill=1)
    
    mask = np.array(mask, dtype=bool)
  
    return mask

# Convert labels shape to array
def shapes_to_label(img_shape, shapes, label_name_to_value, type='class'):    
    assert type in ['class', 'instance']
    mask_shape_array = np.zeros(img_shape[:2], dtype=np.int32)
    
    if type == 'instance':
        ins = np.zeros(img_shape[:2], dtype=np.int32)
        instance_names = ['_background_']
    
    for shape in shapes:
        
        points = shape['points'] # mask points
        label = shape['label'] # mask label name -> String
        shape_type = shape.get('shape_type', None) # polygon or something else
        
        if type == 'class':
            cls_name = label # mask label name -> String
            
        elif type == 'instance':
            cls_name = label.split('-')[0]
            if label not in instance_names:
                instance_names.append(label)
            ins_id = instance_names.index(label)
        
        cls_name = cls_name.strip()
        class_id = label_name_to_value[cls_name] # mask label id in 0, 1, 2...
        mask = shape_to_mask(img_shape[:2], points, shape_type) # mask shape (H * W)
        mask_shape_array[mask] = class_id # set the area mask not zero to mask_id
        
        if type == 'instance':
            ins[mask] = ins_id

    if type == 'instance':
        return mask_shape_array, ins
    
    return mask_shape_array

def json_to_image(img_shape, json_path, toImage=False, visual=False, mask_label='box'):
    """ Covert json labels to segementation masks
    Args:
        img_shape: [h, w, c]
        json_path: json file full name
        toImage:
        visual: whether show the mask on origin image
        mask_label: label name for the mask
    Returns:
        save masks to given output folder
    Raises:
        IOError: An error occurred accessing ablumentation object.
    """
    
    with open(json_path, 'r+', encoding="utf-8") as f:
        """ Load json file and update dictionary order """
        data = json.load(f)

    data['shapes'] = sorted(data['shapes'], key = lambda x: x['label'], reverse=False)
    # imageData = data.get('imageData')
    # # data = json.load(open(os.path.join(path_in, json_file)), encoding="utf-8-sig")
    
    # """ Convert imageData to numpy array """
    # if not imageData:
    #     if Path(data['imagePath']).stem:
    #         img = np.asarray(Image.open(img_path))
    #         # with open(img_path, 'rb') as f:
    #         #     imageData = f.read()
    #         #     imageData = base64.b64encode(imageData).decode('utf-8')
    #         #     img = img_b64_to_arr(imageData)
    #     else:
    #         raise ValueError('Can not find image file for this mask file.')
    # else:
    #     img = img_b64_to_arr(imageData)
    # print(f"01 The dimension of origin image is {img.shape}")

    """ Generate labels in the type of label encoder """
    # label encode each mask class
    mask_to_label_encode = {'_background_': 0}
    for i in range(len(data['shapes'])):
        shape_slice = data['shapes'][i]
        label_slice = shape_slice['label']
        label_slice = label_slice.replace(' ', '')
        if label_slice not in mask_to_label_encode.keys():
            mask_to_label_encode[label_slice] = 1
    # mask_to_label_encode = {'_background_': 0, mask_label: 1} 
    # mask_to_label_encode = {'_background_': 0, '01_bonding': 1, 
    #                         '02_spot': 2, '03_spatter': 3, '04_slag': 4}
    
    
    #         for shape in sorted(data['shapes'], key = lambda x: x['label'], reverse=False): # data['shapes'] -> list, sort by label class
    #             mask_label_name = shape['label'] # Only one shape each round

    #             if mask_label_name in mask_to_label_encode:
    #                 mask_label_value = mask_to_label_encode[mask_label_name] # mask_label_name means 'Joint', 'Line', or something else
    #             else:
    #                 mask_label_value = len(mask_to_label_encode) # If mask not include in list, add new mask label value
    #                 mask_to_label_encode[mask_label_name] = mask_label_value

    """ Generate labels in the type of classID """
    # for s in data['shapes']:
    #     label = s['label']  
    #     label.replace(' ', '')
    #     mask_with_class_id = shapes_to_label(img_shape, [s], mask_to_label_encode) # mask's value are class id
    
    #     yield label, mask_with_class_id
    
    mask_with_class_id = shapes_to_label(img_shape, data['shapes'], mask_to_label_encode) # mask's value are class id
    
    return mask_with_class_id
###################################################################################

def segImage2(img, contours, boundaryColor='yellow', *args):
    #SEGIMAGE Given an image and its segmentation, plot the segment boundaries
    #Input
    # I               input image
    # S              input segmentation
    #varargin   width: width of the boundary
    #                 color: color of the boundary
    # Output
    # I               output image with segment boundaries
    debug = 0
    varargin = args
    if debug:
        t0 = time.time()
    if ( len(varargin) % 2 != 0 ):
            print('Extra Parameters passed to the function must be passed in pairs.')

    parameterCount = int(len(varargin)/2)

    for parameterIndex in range(parameterCount):
            parameterName = varargin[parameterIndex*2 - 1]
            parameterValue = varargin[parameterIndex*2]
            switcher = str.lower(parameterName)
            if switcher ==  'boundarycolor':
                boundaryColor = parameterValue
            elif switcher == 'boundarywidth':
                boundaryWidth = parameterValue
            else:
                print('Invalid parameter')
    
    default_bgr_values = [255, 0, 255]  # ???
    color2value = {'yellow':[0, 255, 255], 'red':[0, 0, 255], 'black':[0, 0, 0], 'pink':[180, 105, 255]}
    boundaryColor = boundaryColor.lower()
    if boundaryColor in color2value:
        bgr_values = color2value[boundaryColor]
    else:
        bgr_values = default_bgr_values
    
    for i in range(len(contours)):
        cv2.drawContours(img, contours, i, bgr_values, 1)
    
    return img

# PNG mask add to img
###################################################################################
def segImage(I, S, boundaryWidth=1, boundaryColor='yellow', *args):
    #SEGIMAGE Given an image and its segmentation, plot the segment boundaries
    #Input
    # I               input image
    # S              input segmentation
    #varargin   width: width of the boundary
    #                 color: color of the boundary
    # Output
    # I               output image with segment boundaries
    debug = 0
    varargin = args
    if debug:
        t0 = time.time()
    if ( len(varargin) % 2 != 0 ):
            print('Extra Parameters passed to the function must be passed in pairs.')

    parameterCount = int(len(varargin)/2)

    for parameterIndex in range(parameterCount):
            parameterName = varargin[parameterIndex*2 - 1]
            parameterValue = varargin[parameterIndex*2]
            switcher = str.lower(parameterName)
            if switcher ==  'boundarycolor':
                boundaryColor = parameterValue
            elif switcher == 'boundarywidth':
                boundaryWidth = parameterValue
            else:
                print('Invalid parameter')
    
    default_bgr_values = [255, 0, 255]  # ???
    color2value = {'yellow':[0, 255, 255], 'red':[0, 0, 255], 'black':[0, 0, 0], 'pink':[180, 105, 255]}
    boundaryColor = boundaryColor.lower()
    if boundaryColor in color2value:
        bgr_values = color2value[boundaryColor]
    else:
        bgr_values = default_bgr_values
    
    # img=cv2.drawContours(img, contours,-1,(0,255,0),5)
    
    S = np.int16(S)
    [cy,cx] = np.gradient(S)
    if debug:
        t1 = time.time()
        print('gradient take time = {}'.format(t1 - t0))
    ccc = np.where((abs(cx)+abs(cy)) >0, 1, 0)
    ccc1 = np.where(ccc > 0, 0, 1)

    ccc = np.uint8(ccc)
    ccc1 = np.uint8(ccc1)
    if boundaryWidth>1:
        boundaryWidth = np.ceil(boundaryWidth)
        dilateWindow = np.ones(boundaryWidth, boundaryWidth)
        ccc = cv2.dilate(ccc, dilateWindow, iterations=1)
    elif boundaryWidth<1:
        print('boundaryWidth has been reset to 1.')
    
    I = I * np.expand_dims(ccc1, axis=2).repeat(3, axis=2)
    for index, layer_value in enumerate(bgr_values):  # BGR
        I[:,:,index] = I[:,:,index] + ccc * layer_value / 255

    if debug:
        t2 = time.time()
        print('segImg take time = {}'.format(t2 - t1))

    return I

def img_add_mask(img_path, mask_path, img_name, code, dst_path, image_quality=95, mask_suffix='', img_code_map=None):
    img = cv2.imread(os.path.join(img_path, code, img_name), 1)
    if (img_code_map is not None) and (img_name not in img_code_map): return
    mask_code = img_code_map[img_name] if(img_code_map is not None and img_name in img_code_map) else code

    suffixs = ['.json', '.png']
    if mask_suffix != '':
        suffixs = [mask_suffix]

    exist_flag = False
    for x in suffixs:
        mask_name = Path(img_name).with_suffix(x)
        mask_p = os.path.join(mask_path, mask_code, mask_name)
        if os.path.exists(mask_p):
            mask_suffix = x
            exist_flag = True
            break
    
    if not exist_flag:
        return 

    if mask_suffix=='.png':
        mask = cv2.imread(mask_p, 0)
        # mask = cv2.dilate(mask, (3,3), iterations=5)
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        # img_with_mask = segImage(img/255, mask, boundaryColor='red')
        img_with_mask = segImage2(img, contours, boundaryColor='red')
    elif mask_suffix=='.json':
        mask = json_to_image(img_shape=img.shape, json_path=mask_p, toImage=False, visual=False, mask_label='box')
        # img_with_mask = segImage(img/255, mask, boundaryColor='red')
        with open(mask_p, 'r+', encoding="utf-8") as f:
            """ Load json file and update dictionary order """
            data = json.load(f)
        
        contours = []
        for s in data['shapes']:
            contours.append(np.expand_dims(s['points'], 1).astype(np.int32))

        img_with_mask = segImage2(img, contours, boundaryColor='red')
        # for label, mask in json_to_image(img_shape=img.shape, 
        #                         json_path=mask_p, toImage=False, visual=False, mask_label='box'):
        #     # label
        #     if (len(label)==6 and label.endswith('1')) or label.lower == 'cut':
        #         boundaryColor = 'red'
        #     elif len(label)==6 and  label.endswith('2'):
        #         boundaryColor = 'black'
        #     elif len(label)==6 and  label.endswith('3'):
        #         boundaryColor = 'pink'
        #     else:
        #         boundaryColor = 'yellow'
        #     img_with_mask = segImage(img_with_mask, mask, boundaryColor=boundaryColor)

    os.makedirs(os.path.join(dst_path, code), exist_ok=True)
    cv2.imwrite(os.path.join(dst_path, code, img_name), img_with_mask, [cv2.IMWRITE_JPEG_QUALITY, image_quality]) 
###################################################################################


# T7 yellow(code) / red(code1) / black(code2) / pink (code3) / 
def main():
    img_path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/T6_Vtech/T665D07/original_data_reshape_20250306'
    mask_path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/T6_Vtech/T665D07/original_data_reshape_20250306'

    # dst_path = img_path + '_ImageWithMask'
    dst_path = img_path + '_ImageWithMask'

    image_quality = 90
    img_suffix = '.jpg'
    mask_suffix='.json'
    codes = []

    df = pd.DataFrame()
    img_name_list = []
    code_list = []
    for code in os.listdir(img_path):
        if (len(codes)!=0) and (code not in codes):
            continue
        imgs = glob.glob(os.path.join(img_path, code, '*'+img_suffix))
        img_name_list.extend([str(Path(x).name) for x in imgs])
        code_list.extend([code] * len(imgs))
    
    df['img'] = img_name_list
    df['code'] = code_list


    df = df.sample(frac=1, random_state=10).reset_index(drop=True)
    
    print('len df: ', len(df))

    # map_path = r"E:\TCL\AutoRepair_T7\data\t7_testing_data\pred_2.csv"
    # data = pd.read_csv(map_path, index_col=False, encoding='ANSI')
    img_code_map =  None
    # for index, row in data.iterrows():
    #     img_code_map[row['img_name']] = row['pred_code']


    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True, nb_workers=5) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False
    
    if is_pandarallel:
        df.parallel_apply(lambda x:img_add_mask(img_path=img_path, mask_path=mask_path, img_name=x['img'], code=x['code'], 
                                                dst_path=dst_path, image_quality=image_quality, mask_suffix=mask_suffix, img_code_map=img_code_map), axis=1)
    else:
        df.apply(lambda x:img_add_mask(img_path=img_path, mask_path=mask_path, img_name=x['img'], code=x['code'], 
                                                dst_path=dst_path, image_quality=image_quality, mask_suffix=mask_suffix, img_code_map=img_code_map), axis=1)
    

if __name__ == '__main__':
    main()
